import vtkmodules.all as vtk
from CloudEngine.core.func.create_lut import create_lut


class VectorActor:
    def __init__(self, input_data, cell_normals, options=None, scale=1.0):
        self.input_data = input_data
        self.cell_normals = cell_normals
        self.options = options or {}
        self.instance = vtk.vtkActor()
        self.scale = scale

        self.mapper = vtk.vtkPolyDataMapper()
        self.create_vector_field()
        self.set_visibility(False)

    def create_vector_field(self):
        arrow = vtk.vtkArrowSource()

        # 创建一个新的 vtkPolyData 对象，用于向量场
        self.vector_data = vtk.vtkPolyData()
        self.vector_data.ShallowCopy(self.input_data)

        self.glyph = vtk.vtkGlyph3D()
        self.glyph.SetSourceConnection(arrow.GetOutputPort())
        self.glyph.SetInputData(self.vector_data)
        self.glyph.SetVectorModeToUseVector()
        self.glyph.SetScaleModeToDataScalingOff()
        self.glyph.SetScaleFactor(self.scale)
        self.glyph.OrientOn()

        # 使用向量的大小作为标量值
        vectors = self.vector_data.GetPointData().GetVectors()
        if vectors:
            magnitudes = vtk.vtkFloatArray()
            magnitudes.SetNumberOfComponents(1)
            magnitudes.SetNumberOfTuples(vectors.GetNumberOfTuples())
            for i in range(vectors.GetNumberOfTuples()):
                vector = vectors.GetTuple3(i)
                magnitude = (vector[0] ** 2 + vector[1] ** 2 + vector[2] ** 2) ** 0.5
                magnitudes.SetValue(i, magnitude)

            self.vector_data.GetPointData().SetScalars(magnitudes)

        self.glyph.Update()

        self.mapper.SetInputConnection(self.glyph.GetOutputPort())
        self.instance.SetMapper(self.mapper)

    def set_color_mapping(self, lut=None, scalar_range=None):

        lut = create_lut(self.options.get('colorType', "Rainbow Desaturated"),
                         self.options.get('colorLevels', 256))
        lut.SetTableRange(scalar_range)

        self.mapper.SetLookupTable(lut)
        self.mapper.SetScalarRange(scalar_range)
        self.mapper.SetScalarVisibility(True)
        self.mapper.Update()

        self.instance.SetMapper(self.mapper)
        self.instance.Modified()

    def set_visibility(self, visible):
        if self.instance:
            self.instance.SetVisibility(visible)

    def get_instance(self):
        return self.instance

    def create_cell_color(self, cell_id=None, color=None):
        if color is None:
            color = [1, 0, 0, 255]
        if isinstance(color[0], float) or isinstance(color[1], float) or isinstance(
                color[2], float):
            color = [int(i * 255) for i in color]

        input_data = self.mapper.GetInput()
        cell_data = input_data.GetCellData()
        colors = cell_data.GetScalars()
        npts = input_data.GetNumberOfCells()

        if colors is None:
            colors = vtk.vtkUnsignedCharArray()
            colors.SetNumberOfComponents(4)
            colors.SetName("colors")

            for i in range(npts):
                colors.InsertNextTuple4(*color)
        if cell_id is not None:
            colors.SetTuple4(cell_id, *color)
        else:
            for i in range(npts):
                colors.SetTuple4(i, *color)

        cell_data.SetScalars(colors)

        input_data.Modified()
